#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
CPA (Correlation Power Analysis) 相关Power AnalysisAttack
"""

import numpy as np
import matplotlib.pyplot as plt
import time
from typing import Dict, Any
try:
    from .base_side_channel_attack import PowerAnalysisAttack, AttackResult
except ImportError:
    from base_side_channel_attack import PowerAnalysisAttack, AttackResult


class CPAAttack(PowerAnalysisAttack):
    """相关Power AnalysisAttack"""
    
    def __init__(self):
        super().__init__("CPA")
        
    def attack(self, target_byte: int = 0, **kwargs) -> AttackResult:
        """
        ExecuteCPAAttack
        
        Args:
            target_byte: 目标字节位置
            **kwargs: 其他Parameters
            
        Returns:
            AttackResult
        """
        start_time = time.time()
        
        if not self.validate_data():
            return AttackResult(
                attack_type="CPA",
                target_byte=target_byte,
                execution_time=0,
                success=False,
                confidence=0.0
            )
        
        try:
            # 获取目标Data
            target_data = self.get_target_data(target_byte)
            
            # CalculateAllFalse设的相关性
            correlations = np.zeros((self.hypothesis_range, self.traces.shape[1]))
            
            for hypothesis in range(self.hypothesis_range):
                # CalculateIntermediate value（汉明权重）
                intermediate_values = self.calculate_intermediate_values(target_data, hypothesis)
                hamming_weights = self._calculate_hamming_weight(intermediate_values)
                
                # CalculateEachSampling点的相关性
                for sample_point in range(self.traces.shape[1]):
                    correlations[hypothesis, sample_point] = self.calculate_correlation(
                        hamming_weights, sample_point
                    )
            
            # 找To最大相关性
            max_corr_indices = np.unravel_index(
                np.argmax(np.abs(correlations)), correlations.shape
            )
            best_hypothesis = max_corr_indices[0]
            best_sample_point = max_corr_indices[1]
            max_correlation = correlations[best_hypothesis, best_sample_point]
            
            # CalculateEachFalse设的最大相关性（跨AllSampling点）
            max_corr_per_hypothesis = np.max(np.abs(correlations), axis=1)
            
            # 获取前 n 个候选（Default前 256 个，i.e.All）
            top_n = kwargs.get('top_n', min(256, self.hypothesis_range))
            sorted_indices = np.argsort(max_corr_per_hypothesis)[::-1]  # 降序排列
            top_candidates = [
                (int(idx), float(max_corr_per_hypothesis[idx])) 
                for idx in sorted_indices[:top_n]
            ]
            
            execution_time = time.time() - start_time
            
            # Calculate置信度
            confidence = abs(max_correlation)
            success = confidence > 0.1  # 阈值可调
            
            self.logger.info(f"CPA Attack completed:")
            self.logger.info(f"  - Best hypothesis: 0x{best_hypothesis:02x}")
            self.logger.info(f"  - Max correlation: {max_correlation:.4f}")
            self.logger.info(f"  - Sample point: {best_sample_point}")
            self.logger.info(f"  - Confidence: {confidence:.4f}")
            self.logger.info(f"  - Top candidates saved: {len(top_candidates)}")
            
            return AttackResult(
                attack_type="CPA",
                target_byte=target_byte,
                execution_time=execution_time,
                success=success,
                confidence=confidence,
                best_hypothesis=best_hypothesis,
                top_candidates=top_candidates,
                metadata={
                    'correlations': correlations,
                    'max_correlation': max_correlation,
                    'best_sample_point': best_sample_point,
                    'correlation_threshold': 0.1,
                    'max_corr_per_hypothesis': max_corr_per_hypothesis
                }
            )
            
        except Exception as e:
            self.logger.error(f"CPA attack failed: {e}")
            return AttackResult(
                attack_type="CPA",
                target_byte=target_byte,
                execution_time=time.time() - start_time,
                success=False,
                confidence=0.0
            )
    
    def _calculate_hamming_weight(self, values: np.ndarray) -> np.ndarray:
        """
        Calculate汉明权重
        
        Note：Ifvalues已经Yes汉明权重（Float数），直接Return
        """
        # CheckYesNo已经Yes汉明权重（Float数Or小于8的Integer）
        if values.dtype == np.float64 or (values.dtype in [np.int64, np.int32] and np.all(values <= 8)):
            # 已经Yes汉明权重，直接Return
            return values.astype(np.float64)
        
        # No则Calculate汉明权重
        return np.array([bin(int(val)).count('1') for val in values])
    
    def plot_results(self, result: AttackResult) -> None:
        """绘制CPAAttackResult"""
        if not result.success or 'correlations' not in result.metadata:
            self.logger.warning("No valid results to plot")
            return
        
        correlations = result.metadata['correlations']
        
        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 8))
        
        # 绘制相关性热图
        im1 = ax1.imshow(np.abs(correlations), aspect='auto', cmap='hot', interpolation='nearest')
        ax1.set_title('CPA Attack - Correlation Heatmap')
        ax1.set_xlabel('Sample Points')
        ax1.set_ylabel('Key Hypothesis')
        plt.colorbar(im1, ax=ax1, label='|Correlation|')
        
        # 绘制最佳False设的相关性曲线
        best_hypothesis = result.best_hypothesis
        ax2.plot(np.abs(correlations[best_hypothesis, :]), 'b-', linewidth=2)
        ax2.set_title(f'Best Hypothesis (0x{best_hypothesis:02x}) Correlation Trace')
        ax2.set_xlabel('Sample Points')
        ax2.set_ylabel('|Correlation|')
        ax2.grid(True, alpha=0.3)
        
        # 标记最大相关性点
        best_sample_point = result.metadata['best_sample_point']
        max_correlation = result.metadata['max_correlation']
        ax2.plot(best_sample_point, abs(max_correlation), 'ro', markersize=8, 
                label=f'Max: {abs(max_correlation):.4f}')
        ax2.legend()
        
        plt.tight_layout()
        plt.show()
        
        # 绘制AllFalse设的最大相关性
        fig2, ax3 = plt.subplots(1, 1, figsize=(10, 6))
        max_correlations_per_hypothesis = np.max(np.abs(correlations), axis=1)
        ax3.bar(range(self.hypothesis_range), max_correlations_per_hypothesis, 
               color='skyblue', alpha=0.7)
        ax3.bar(best_hypothesis, max_correlations_per_hypothesis[best_hypothesis], 
               color='red', alpha=0.8, label=f'Best: 0x{best_hypothesis:02x}')
        ax3.set_title('Maximum Correlation per Key Hypothesis')
        ax3.set_xlabel('Key Hypothesis')
        ax3.set_ylabel('Max |Correlation|')
        ax3.legend()
        ax3.grid(True, alpha=0.3)
        
        plt.tight_layout()
        plt.show()
